library(maptools)
library(tidyverse)
library(sp)
library(sf)
library(raster)
library(rgdal)
library(PNWColors)
library(MetBrewer)
library(viridis)
library(fixest)


#define some useful functions
add.alpha <- function(col, alpha=1){
  if(missing(col))
    stop("Please provide a vector of colours.")
  apply(sapply(col, col2rgb)/255, 2, 
        function(x) 
          rgb(x[1], x[2], x[3], alpha=alpha))  
}	
get_grid_raster <- function(raster){
  output <- raster(raster)
  output[] <- 1:ncell(output)
  return(output)
}
identify_non_cells <- function(raster, poly){
  raster = raster(raster)
  raster[] <- 1:ncell(raster)
  r_ext <- extract(raster, poly)
  all_cells <- 1:ncell(raster)
  non_cells <- all_cells[all_cells %in% unlist(r_ext) == F]
  return(non_cells)
  
}




        #load publicly available raw data and prep for plotting
    
        #       ############
        #       
        #       ## PM2.5
        # 
        #         grid <- raster(raster("1_input/data/vandonk/annual_001/V5GL02.HybridPM25.Global.199801-199812.nc"))
        #         grid[] <- 1:ncell(grid)
        #         
        #         
        #         fls <- list.files("1_input/data/vandonk/annual_001/")[13:22][6:10]
        #         
        #         
        #         rdat <- raster::brick(x = grid, nl = length(fls))
        #         
        #         for (i in 1:length(fls)){
        #           
        #           rdat[[i]] <- (raster(paste("1_input/data/vandonk/annual_001/", fls[i], sep = "")))
        #           
        #         }
        #         
        #         rdat <- raster::stackApply(x=rdat,indices = 1, fun= mean)
        #         
        #         pm <- rdat
        #         rm(rdat)
        #         
        #      
        # ############################################
        #         
        #        ## Wealth ## 
        #         
        #         wfls <- list.files("1_input/data/rwi/")
        #         wfls_cty <- substr(wfls,1,3)
        #         
        #         
        #         rwi_list <- list()
        #         
        #         for (i in 1:length(wfls)){
        #           rwi_list[[i]] <- read_csv(paste("1_input/data/rwi/",wfls[i],sep=""))
        #           rwi_list[[i]]$country <- wfls_cty[i]
        #         }      
        #         
        #         
        #         rwi <- data.frame(data.table::rbindlist(rwi_list))
        #         
        #         
        # ############################################
        #         
        #         
        #   ur <- raster("1_input/data/GIS/rasters/Urban-Rural Catchment Areas (URCA).tif")
        #         
        # 
        #         
        # #### Combine ######
        #         
        #         
        #         
        #   data <- rwi; rm(rwi)
        #   data$pm <- pm[cellFromXY(pm, as.matrix(data[,c("longitude","latitude")]))]
        #   data$urca <- ur[cellFromXY(ur, as.matrix(data[,c("longitude","latitude")]))]
        #   data$urban <- as.numeric(data$urca %in% 1:7)
        #   data$rural <- as.numeric(data$urca %in% 15:30)
        #   data$peri <- as.numeric(data$urca %in% 8:14)
        #   data <- data %>% dplyr::select(country, lon = longitude, lat = latitude, rwi, error, pm, urca, urban, rural)
        #   data$cellpm <- cellFromXY(pm, as.matrix(data[,c("longitude","latitude")]))
        #   
        #   
        # 
        #   data_mat <- data
        #   data_mat$pm[data_mat$pm>100]<-100
        #   data_mat$rwi[data_mat$rwi< -1]<- -1
        #   data_mat$rwi[data_mat$rwi>1]<-1
        #   
        #   
        #   
        #   
        #   
        #   ##[a]
        #         mat <- matrix(nrow =length(seq(0,110,1))-1, ncol = length(seq(-1, 1.7, .05)) )
        #         for(i in 1:(nrow(mat) -1)){
        #           for(j in 1:(ncol(mat)-1 ) ){
        #             
        #             mat[i,j]<-sum(
        #               data$rwi > seq(-1, 1.7, .05)[j] & data$rwi <= seq(-1, 1.7, .05)[j+1] &
        #                 data$pm > (seq(0,110,1))[i] & data$pm <=  (seq(0,110,1))[i+1]
        #               , na.rm =T)
        #           }
        #         }
        #         #mat[mat==0]<-NA
        #         mat1 = mat
        #         col1 <- apply(mat1, 1, function(x){sum(x, na.rm = T)})
        #         row1 <- apply(mat1, 2, function(x){sum(x, na.rm = T)})
        #         
        #         
        #   ##[b]
        #         
        #         mat <- matrix(nrow =length(seq(0,110,1))-1, ncol = length(seq(-1, 1.7, .05)) )
        #         for(i in 1:(nrow(mat) -1)){
        #           for(j in 1:(ncol(mat)-1 ) ){
        #             
        #             mat[i,j]<-sum(
        #               data$rwi[data$urban == 1] > seq(-1, 1.7, .05)[j] & data$rwi[data$urban == 1] <= seq(-1, 1.7, .05)[j+1] &
        #                 data$pm[data$urban == 1] > (seq(0,110,1))[i] & data$pm[data$urban == 1] <=  (seq(0,110,1))[i+1]
        #               , na.rm =T)
        #           }
        #         }
        #         #mat[mat==0]<-NA
        #         mat2 = mat
        #         col2 <- apply(mat2, 1, function(x){sum(x, na.rm = T)})
        #         row2 <- apply(mat2, 2, function(x){sum(x, na.rm = T)})
        #      
        #         
        #         
        #            
        #   ##[c] 
        #         mat <- matrix(nrow =length(seq(0,110,1))-1, ncol = length(seq(-1, 1.7, .05)) )
        #         for(i in 1:(nrow(mat) -1)){
        #           for(j in 1:(ncol(mat)-1 ) ){
        #             
        #             mat[i,j]<-sum(
        #               data$rwi[data$rural == 1] > seq(-1, 1.7, .05)[j] & data$rwi[data$rural == 1] <= seq(-1, 1.7, .05)[j+1] &
        #                 data$pm[data$rural == 1] > (seq(0,110,1))[i] & data$pm[data$rural == 1] <=  (seq(0,110,1))[i+1]
        #               , na.rm =T)
        #           }
        #         }
        #         #mat[mat==0]<-NA
        #         mat3 = mat
        #         col3 <- apply(mat3, 1, function(x){sum(x, na.rm = T)})
        #         row3 <- apply(mat3, 2, function(x){sum(x, na.rm = T)})
        #         
        #         
        #         
        #       #[d] periurban  
        #      
        #        mat <- matrix(nrow =length(seq(0,110,1))-1, ncol = length(seq(-1, 1.7, .05)) )
        #         for(i in 1:(nrow(mat) -1)){
        #           for(j in 1:(ncol(mat)-1 ) ){
        #             
        #             mat[i,j]<-sum(
        #               data$rwi[data$peri == 1] > seq(-1, 1.7, .05)[j] & data$rwi[data$peri == 1] <= seq(-1, 1.7, .05)[j+1] &
        #                 data$pm[data$peri == 1] > (seq(0,110,1))[i] & data$pm[data$peri == 1] <=  (seq(0,110,1))[i+1]
        #               , na.rm =T)
        #           }
        #         }
        #         #mat[mat==0]<-NA
        #         mat4 = mat
        #         col4 <- apply(mat4, 1, function(x){sum(x, na.rm = T)})
        #         row4 <- apply(mat4, 2, function(x){sum(x, na.rm = T)})
        #         
        #         
          
    load("data/figSI2-plotdata.RData")
  
  
        
        
  pdf(file = "figures/fig_SI2.pdf", width = 16, height = 4)
  
      par(mar = c(4,3,1,0))
      par(mgp = c(3,1.75,0))
      par(mfrow =c(1,4))
      
      
      #[a] full sample
      mat = mat1
      row = row1
      col = col1
      
      plot(1, 1, col = NA, xlab = "",ylab = "", axes = F,  xlim = c(-0.25, 1), ylim = c(-0.25, 1))
      plot(flip(raster(mat), direction= 'y'), col = pnw_palette("Winter",100)[100:1], legend = F, box = F, add=T)
      axis(1, at = seq(0,1,1/5), labels = seq(-1, 1.7, .5))
      axis(2, at = seq(0,100/110,(100/110)/(length(seq(0, 100, 20))-1)), labels = seq(0, 100, 20), las = 2)
      
      x = seq(0,1,1/54)
      y = .2*row/max(row) - .25
      lines(seq(0,1,1/54), .2*row/max(row) - .25, lwd = 2, col = 'gray50')
      polygon(c( x, rep(0, length(x))),c( y, rev(y)), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      x = seq(0,1,1/109)
      y = .2*col/max(col) - .25
      lines( .2*col/max(col) - .25,seq(0,1,1/109), lwd = 2, col = 'gray50')
      polygon(c( y, rev(y)),c( x, rep(0, length(x))), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      
      #[b] rurla
      mat = mat3
      row = row3
      col = col3
      
      
      
      plot(1, 1, col = NA, xlab = "",ylab = "", axes = F,  xlim = c(-0.25, 1), ylim = c(-0.25, 1))
      plot(flip(raster(mat), direction= 'y'), col = pnw_palette("Winter",100)[100:1], legend = F, box = F, add=T)
      axis(1, at = seq(0,1,1/5), labels = seq(-1, 1.7, .5))
      
      x = seq(0,1,1/54)
      y = .2*row/max(row) - .25
      lines(seq(0,1,1/54), .2*row/max(row) - .25, lwd = 2, col = 'gray50')
      polygon(c( x, rep(0, length(x))),c( y, rev(y)), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      x = seq(0,1,1/109)
      y = .2*col/max(col) - .25
      lines( .2*col/max(col) - .25,seq(0,1,1/109), lwd = 2, col = 'gray50')
      polygon(c( y, rev(y)),c( x, rep(0, length(x))), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
    
      
      
      #[c] peri-urban
      mat = mat4
      row = row4
      col = col4
      
      plot(1, 1, col = NA, xlab = "",ylab = "", axes = F,  xlim = c(-0.25, 1), ylim = c(-0.25, 1))
      plot(flip(raster(mat), direction= 'y'), col = pnw_palette("Winter",100)[100:1], legend = F, box = F, add=T)
      axis(1, at = seq(0,1,1/5), labels = seq(-1, 1.7, .5))
      axis(2, at = seq(0,100/110,(100/110)/(length(seq(0, 100, 20))-1)), labels = seq(0, 100, 20), las = 2)
      
      x = seq(0,1,1/54)
      y = .2*row/max(row) - .25
      lines(seq(0,1,1/54), .2*row/max(row) - .25, lwd = 2, col = 'gray50')
      polygon(c( x, rep(0, length(x))),c( y, rev(y)), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      x = seq(0,1,1/109)
      y = .2*col/max(col) - .25
      lines( .2*col/max(col) - .25,seq(0,1,1/109), lwd = 2, col = 'gray50')
      polygon(c( y, rev(y)),c( x, rep(0, length(x))), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      
      
      #[d] urban
      mat = mat2
      row = row2
      col = col2
     
      plot(1, 1, col = NA, xlab = "",ylab = "", axes = F,  xlim = c(-0.25, 1), ylim = c(-0.25, 1))
      plot(flip(raster(mat), direction= 'y'), col = pnw_palette("Winter",100)[100:1], legend = F, box = F, add=T)
      axis(1, at = seq(0,1,1/5), labels = seq(-1, 1.7, .5))
      
      x = seq(0,1,1/54)
      y = .2*row/max(row) - .25
      lines(seq(0,1,1/54), .2*row/max(row) - .25, lwd = 2, col = 'gray50')
      polygon(c( x, rep(0, length(x))),c( y, rev(y)), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
      x = seq(0,1,1/109)
      y = .2*col/max(col) - .25
      lines( .2*col/max(col) - .25,seq(0,1,1/109), lwd = 2, col = 'gray50')
      polygon(c( y, rev(y)),c( x, rep(0, length(x))), col = add.alpha(pnw_palette("Winter",100)[50], .5) , border = NA)  
      
      
  
  
  
     
      
    dev.off()
  
  
  
    
    
    
    
  
  
  
